package com.krafteapps.androidsecurehttpclient;

import java.io.File;
import java.io.FileInputStream;
import java.io.FileNotFoundException;
import java.io.IOException;
import java.security.KeyStore;
import java.security.KeyStoreException;
import java.security.NoSuchAlgorithmException;
import java.security.cert.CertificateException;
import java.security.cert.X509Certificate;
import java.util.HashMap;
import java.util.Map;

import javax.net.ssl.TrustManager;
import javax.net.ssl.TrustManagerFactory;
import javax.net.ssl.X509TrustManager;

import android.content.Context;
import android.util.Log;

public class HttpTrustManager {
	private static final String TAG = "HttpTrustManager";
	private static Map<File, HttpTrustManager> mTrustManagerMap = new HashMap<File, HttpTrustManager>();
	
	private int mInstanceCount;
	private File mKeyStoreFile;
	private KeyStore mKeyStore;

	private X509TrustManager mDefaultTrustManager;
	private X509TrustManager mUnsecureTrustManager;
	private X509TrustManager mLocalTrustManager;
	
	private HttpTrustManager(File keyStoreFile) {
		mInstanceCount = 1;
		mKeyStoreFile = keyStoreFile;
	}

	protected static HttpTrustManager getInstance(File keyStoreFile) {
		HttpTrustManager tm;
		synchronized (mTrustManagerMap) {
			if (mTrustManagerMap.containsKey(keyStoreFile)) {
				tm = mTrustManagerMap.get(keyStoreFile);
				tm.mInstanceCount++;
			} else {
				tm = new HttpTrustManager(keyStoreFile);
				mTrustManagerMap.put(keyStoreFile, tm);
			}
		}
		
		synchronized (tm) {
			try {
				TrustManagerFactory tmf;
				TrustManager[] tms;
				
				tm.mKeyStore = KeyStore.getInstance(KeyStore.getDefaultType());
				
				FileInputStream fis;
				try {
					fis = new FileInputStream(keyStoreFile);
				} catch (FileNotFoundException e) {
					fis = null;
				}
				
				try {
					tm.mKeyStore.load(fis, "".toCharArray());
				} catch (CertificateException e) {
					e.printStackTrace();
					tm.mKeyStore = null;
				} catch (IOException e) {
					e.printStackTrace();
					tm.mKeyStore = null;
				}
				
				// Load local trust manager based on keystore
				tmf = TrustManagerFactory.getInstance("X509");
				tmf.init(tm.mKeyStore);
				tms = tmf.getTrustManagers();
				if (tms!=null) {
					for (TrustManager t : tms) {
	                    if (t instanceof X509TrustManager) {
	                        tm.mLocalTrustManager = (X509TrustManager) t;
	                        break;
	                    }
	                }
				}
				
				// Load default trust manager
				 tmf = TrustManagerFactory.getInstance("X509");
		         tmf.init((KeyStore)null);
		         tms = tmf.getTrustManagers();
		         if (tms != null) {
		        	 for (TrustManager t : tms) {
	                    if (t instanceof X509TrustManager) {
	                    	tm.mDefaultTrustManager = (X509TrustManager) t;
	                    	break;
	                    }
		        	 }
		         }
				
			} catch (NoSuchAlgorithmException e) {
				// TODO Auto-generated catch block
				e.printStackTrace();
			} catch (KeyStoreException e) {
				//TODO Auto-generated catch block
				e.printStackTrace();
			}
			
			tm.mUnsecureTrustManager = new SimpleX509TrustManager();
		}
		
		return tm;
	}
	
	protected void releaseInstance() {
		synchronized (mTrustManagerMap) {
			if (mTrustManagerMap.containsKey(mKeyStoreFile)) {
				mInstanceCount--;
				if (mInstanceCount <= 0)
					mTrustManagerMap.remove(mKeyStoreFile);
			}
		}
	}

	protected KeyStore getKeyStore() {
		return mKeyStore;
	}
	
	public static void clearKeyStore(Context context) {
		clearKeyStore(AndroidSecureHttpClient.getDefaultKeyStoreFile(context));
	}
	
	public static void clearKeyStore(File file) {
		file.delete();
		synchronized (mTrustManagerMap) {
			if (mTrustManagerMap.containsKey(file)) {
				// KeyStore was loaded in memory, update
				HttpTrustManager tm = mTrustManagerMap.get(file);
				
				TrustManagerFactory tmf;
				TrustManager[] tms;
				
				try {
					tm.mKeyStore = KeyStore.getInstance(KeyStore.getDefaultType());
					
					try {
						tm.mKeyStore.load(null, "".toCharArray());
					} catch (CertificateException e) {
						e.printStackTrace();
						tm.mKeyStore = null;
					} catch (IOException e) {
						e.printStackTrace();
						tm.mKeyStore = null;
					}
					
					// Load local trust manager based on keystore
					tmf = TrustManagerFactory.getInstance("X509");
					tmf.init(tm.mKeyStore);
					tms = tmf.getTrustManagers();
					if (tms!=null) {
						for (TrustManager t : tms) {
		                    if (t instanceof X509TrustManager) {
		                        tm.mLocalTrustManager = (X509TrustManager) t;
		                        break;
		                    }
		                }
					}
				} catch (NoSuchAlgorithmException e) {
					//TODO Auto-generated catch block
					e.printStackTrace();
				} catch (KeyStoreException e) {
					//TODO Auto-generated catch block
					e.printStackTrace();
				}
			}
		}
	}
	
	public void addCertificateChain(X509Certificate[] chain) throws CertificateException {
		try {
            javax.net.ssl.TrustManagerFactory tmf = javax.net.ssl.TrustManagerFactory.getInstance("X509");
            for (X509Certificate element : chain)
                mKeyStore.setCertificateEntry(element.getSubjectDN().toString(), element);

            tmf.init(mKeyStore);
            TrustManager[] tms = tmf.getTrustManagers();
            if (tms != null) {
                for (TrustManager tm : tms) {
                    if (tm instanceof X509TrustManager) {
                        mLocalTrustManager = (X509TrustManager) tm;
                        break;
                    }
                }
            }
            java.io.FileOutputStream keyStoreStream;
            try {
                keyStoreStream = new java.io.FileOutputStream(mKeyStoreFile);
                mKeyStore.store(keyStoreStream, "".toCharArray());
                keyStoreStream.close();
            } catch (FileNotFoundException e) {
                throw new CertificateException("Unable to write KeyStore: " + e.getMessage());
            } catch (CertificateException e) {
                throw new CertificateException("Unable to write KeyStore: " + e.getMessage());
            } catch (IOException e) {
                throw new CertificateException("Unable to write KeyStore: " + e.getMessage());
            }

        } catch (NoSuchAlgorithmException e) {
            Log.e(TAG, "Unable to get X509 Trust Manager ", e);
        } catch (KeyStoreException e) {
            Log.e(TAG, "Key Store exception while initializing TrustManagerFactory ", e);
        }
	}

	protected TrustManager getNewTrustManager(boolean secure) {
		if (secure)
			return new SecureX509TrustManager();
		else
			return mUnsecureTrustManager;
	}
	
	private static class SimpleX509TrustManager implements X509TrustManager {
        public void checkClientTrusted(X509Certificate[] chain, String authType) throws CertificateException {
        }

        public void checkServerTrusted(X509Certificate[] chain, String authType) throws CertificateException {
        }

        public X509Certificate[] getAcceptedIssuers() {
            return null;
        }
    }

    private class SecureX509TrustManager implements X509TrustManager {
        private SecureX509TrustManager() {
        }

        public void checkClientTrusted(X509Certificate[] chain, String authType) throws CertificateException {
            mDefaultTrustManager.checkClientTrusted(chain, authType);
        }

        public void checkServerTrusted(X509Certificate[] chain, String authType) throws CertificateNotAcceptedException, CertificateException {
            try {
                mDefaultTrustManager.checkServerTrusted(chain, authType);
            } catch (CertificateException e) {
            	try {
            		mLocalTrustManager.checkServerTrusted(new X509Certificate[] {chain[0]}, authType);
            	} catch (CertificateException ce) {
            		throw new CertificateNotAcceptedException(chain);
				}
            }
        }

        public X509Certificate[] getAcceptedIssuers()
        {
            return mDefaultTrustManager.getAcceptedIssuers();
        }

    }
    
    @SuppressWarnings("serial")
	public static class CertificateNotAcceptedException extends CertificateException {
    	private static X509Certificate[] mChain;
    	public CertificateNotAcceptedException(X509Certificate[] chain) {
    		super();
    		mChain = chain;
    	}
    	public static X509Certificate[] getLastChain() {
    		return mChain;
    	}
    }
}
